import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

from visualize_knn_boundary import visualize_knn_boundary

'''
        任务：鸢尾花识别
'''

DATA_FILE = './data/Iris.csv'

SPECIES_LABEL_DICT = {
    'Iris-setosa': 0, # 山鸢尾
    'Iris-versicolor': 1, # 变色鸢尾
    'Iris-virginica': 2 # 维吉尼亚鸢尾
}

# 使用的特征列
FEAT_COLS = ['SepalLengthCm', 'SepalWidthCm', 'PetalLengthCm', 'PetalWidthCm']


def investigate_knn(iris_data, sel_cols, k):
    '''
    不同的K值对模型的影响
    '''
    # 获取数据集特征
    X = iris_data[sel_cols].values
    # 获取数据集标签
    y = iris_data['Label'].values
    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=1 / 3, random_state=10)
    # 声明变量
    knn_model = KNeighborsClassifier(n_neighbors=k)

    # 2. 训练模型
    knn_model.fit(X_train, y_train)
    # 3. 模型评估
    score = knn_model.score(X_test, y_test)
    print('模型准确率：{:.2f}%'.format(score * 100))

    # 可视化分类决策边界，选择特征 0（花萼长度）和特征 2（花瓣长度）
    visualize_knn_boundary(knn_model, X_train, y_train, X_test, y_test, feature_indices=(0, 1))


def main():
    # 1. 读取数据
    iris_data = pd.read_csv(DATA_FILE, index_col='Id')
    iris_data['Label'] = iris_data['Species'].map(SPECIES_LABEL_DICT)
    
    k_values = [1, 3, 5, 10, 15]
    sel_cols = ['SepalLengthCm', 'PetalLengthCm']
    for k in k_values:
        print('KNN算法，K值：{}，特征列：{}'.format(k, sel_cols))
        investigate_knn(iris_data, sel_cols,k)




if __name__ == '__main__':
    main()